//! 符号表实现
//! 
//! 管理标签定义和引用，支持地址计算和两遍扫描

use std::collections::HashMap;
use crate::error::{AssemblerError, Result};

/// 符号类型
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum SymbolType {
    /// 标签（代码地址）
    Label,
}

/// 符号信息
#[derive(Debug, Clone)]
pub struct Symbol {
    /// 符号名称
    pub name: String,
    /// 符号类型
    pub symbol_type: SymbolType,
    /// 符号地址（字节偏移）- 使用 u64 支持所有架构
    pub address: u64,
    /// 是否已定义
    pub defined: bool,
    /// 定义位置（行号）
    pub line: usize,
}

/// 符号表
pub struct SymbolTable {
    /// 符号映射表：名称 -> 符号信息
    symbols: HashMap<String, Symbol>,
    /// 当前代码地址（字节偏移）- 使用 u64 支持所有架构
    current_address: u64,
}

impl SymbolTable {
    /// 创建新的符号表
    pub fn new() -> Self {
        Self {
            symbols: HashMap::new(),
            current_address: 0,
        }
    }

    /// 定义标签
    /// 
    /// # 参数
    /// - `name`: 标签名称
    /// - `line`: 行号
    /// - `allow_redefinition`: 是否允许重新定义（第二遍扫描时为 true，用于更新地址）
    pub fn define_label(&mut self, name: String, line: usize, allow_redefinition: bool) -> Result<()> {
        // 检查是否重复定义
        if let Some(existing) = self.symbols.get(&name) {
            if existing.defined {
                if allow_redefinition {
                    // 第二遍扫描：更新地址
                    if let Some(symbol) = self.symbols.get_mut(&name) {
                        symbol.address = self.current_address;
                    }
                    return Ok(());
                } else {
                    // 第一遍扫描：重复定义是错误
                    return Err(AssemblerError::CodeGenError {
                        message: format!(
                            "标签 '{}' 重复定义（第 {} 行和第 {} 行）",
                            name, existing.line, line
                        ),
                    });
                }
            }
        }

        let symbol = Symbol {
            name: name.clone(),
            symbol_type: SymbolType::Label,
            address: self.current_address,
            defined: true,
            line,
        };

        self.symbols.insert(name, symbol);
        Ok(())
    }

    /// 引用标签（前向引用）
    pub fn reference_label(&mut self, name: String, line: usize) -> Result<()> {
        if !self.symbols.contains_key(&name) {
            // 创建未定义的符号（前向引用）
            let symbol = Symbol {
                name: name.clone(),
                symbol_type: SymbolType::Label,
                address: 0, // 稍后填充
                defined: false,
                line,
            };
            self.symbols.insert(name, symbol);
        }
        Ok(())
    }

    /// 获取标签地址
    pub fn get_label_address(&self, name: &str) -> Result<u64> {
        self.symbols
            .get(name)
            .ok_or_else(|| AssemblerError::CodeGenError {
                message: format!("未定义的标签: {}", name),
            })
            .and_then(|sym| {
                if sym.defined {
                    Ok(sym.address)
                } else {
                    Err(AssemblerError::CodeGenError {
                        message: format!("标签 '{}' 未定义", name),
                    })
                }
            })
    }

    /// 更新当前地址
    pub fn advance_address(&mut self, bytes: usize) {
        self.current_address += bytes as u64;
    }

    /// 获取当前地址
    pub fn current_address(&self) -> u64 {
        self.current_address
    }

    /// 重置地址计数器（用于第二遍扫描）
    pub fn reset_address(&mut self) {
        self.current_address = 0;
    }

    /// 检查是否有未定义的符号
    pub fn check_undefined_symbols(&self) -> Result<()> {
        for symbol in self.symbols.values() {
            if !symbol.defined {
                return Err(AssemblerError::CodeGenError {
                    message: format!("标签 '{}' 未定义（第 {} 行引用）", 
                                   symbol.name, symbol.line),
                });
            }
        }
        Ok(())
    }
}

impl Default for SymbolTable {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_new_symbol_table() {
        let table = SymbolTable::new();
        assert_eq!(table.current_address(), 0);
    }

    #[test]
    fn test_define_label() {
        let mut table = SymbolTable::new();
        table.define_label("开始".to_string(), 1, false).unwrap();
        assert_eq!(table.get_label_address("开始").unwrap(), 0);
    }

    #[test]
    fn test_advance_address() {
        let mut table = SymbolTable::new();
        table.advance_address(10);
        assert_eq!(table.current_address(), 10);
        
        table.define_label("标签".to_string(), 1, false).unwrap();
        assert_eq!(table.get_label_address("标签").unwrap(), 10);
    }

    #[test]
    fn test_forward_reference() {
        let mut table = SymbolTable::new();
        // 先引用标签
        table.reference_label("标签".to_string(), 1).unwrap();
        // 后定义标签
        table.advance_address(5);
        table.define_label("标签".to_string(), 2, false).unwrap();
        assert_eq!(table.get_label_address("标签").unwrap(), 5);
    }

    #[test]
    fn test_duplicate_label_error() {
        let mut table = SymbolTable::new();
        table.define_label("标签".to_string(), 1, false).unwrap();
        // 第一遍扫描时重复定义应该报错
        let result = table.define_label("标签".to_string(), 2, false);
        assert!(result.is_err());
    }

    #[test]
    fn test_allow_redefinition_second_pass() {
        let mut table = SymbolTable::new();
        table.define_label("标签".to_string(), 1, false).unwrap();
        table.advance_address(10);
        // 第二遍扫描时允许更新地址
        table.define_label("标签".to_string(), 1, true).unwrap();
        assert_eq!(table.get_label_address("标签").unwrap(), 10);
    }

    #[test]
    fn test_undefined_label() {
        let mut table = SymbolTable::new();
        table.reference_label("未定义标签".to_string(), 1).unwrap();
        let result = table.get_label_address("未定义标签");
        assert!(result.is_err());
    }

    #[test]
    fn test_check_undefined_symbols() {
        let mut table = SymbolTable::new();
        table.reference_label("未定义".to_string(), 1).unwrap();
        let result = table.check_undefined_symbols();
        assert!(result.is_err());
    }

    #[test]
    fn test_reset_address() {
        let mut table = SymbolTable::new();
        table.advance_address(100);
        assert_eq!(table.current_address(), 100);
        table.reset_address();
        assert_eq!(table.current_address(), 0);
    }

    #[test]
    fn test_multiple_labels() {
        let mut table = SymbolTable::new();
        table.define_label("标签1".to_string(), 1, false).unwrap();
        table.advance_address(10);
        table.define_label("标签2".to_string(), 2, false).unwrap();
        table.advance_address(5);
        table.define_label("标签3".to_string(), 3, false).unwrap();
        
        assert_eq!(table.get_label_address("标签1").unwrap(), 0);
        assert_eq!(table.get_label_address("标签2").unwrap(), 10);
        assert_eq!(table.get_label_address("标签3").unwrap(), 15);
    }
}

